import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import sys
import os

# 添加模块路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from moudle import (
    create_data_loaders, 
    Trainer, 
    setup_device, 
    load_label_mapping,
    TensorBoardLogger,
    YAMLConfig,
    load_model
)


class BasicBlock(nn.Module):
    """ResNet基础块（用于ResNet18和ResNet34）"""
    expansion = 1
    
    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion * planes)
            )
    
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class Bottleneck(nn.Module):
    """ResNet瓶颈块（用于ResNet50、ResNet101和ResNet152）"""
    expansion = 4
    
    def __init__(self, in_planes, planes, stride=1):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion * planes)
            )
    
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    """ResNet模型实现"""
    
    def __init__(self, block, num_blocks, num_classes=100, dropout=0.0):
        super(ResNet, self).__init__()
        self.in_planes = 64
        
        # 第一层卷积
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        # ResNet层
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        
        # 全局平均池化和分类器
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        self.fc = nn.Linear(512 * block.expansion, num_classes)
        
        # 权重初始化
        self._initialize_weights()
    
    def _make_layer(self, block, planes, num_blocks, stride):
        """构建ResNet层"""
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)
    
    def _initialize_weights(self):
        """初始化权重"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
        
        # 对于Bottleneck块，将最后一个BN层的权重初始化为0
        for m in self.modules():
            if isinstance(m, Bottleneck):
                nn.init.constant_(m.bn3.weight, 0)
            elif isinstance(m, BasicBlock):
                nn.init.constant_(m.bn2.weight, 0)
    
    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.fc(x)
        
        return x


def ResNet18(num_classes=100):
    """ResNet18模型"""
    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes)


def ResNet34(num_classes=100):
    """ResNet34模型"""
    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes)


def ResNet50(num_classes=100, dropout=0.0):
    """ResNet50模型"""
    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes, dropout)


def ResNet101(num_classes=100):
    """ResNet101模型"""
    return ResNet(Bottleneck, [3, 4, 23, 3], num_classes)


def ResNet152(num_classes=100):
    """ResNet152模型"""
    return ResNet(Bottleneck, [3, 8, 36, 3], num_classes)


class ResNetTrainer(Trainer):
    """ResNet专用训练器，可以添加一些ResNet特有的功能"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 可以添加ResNet特有的训练策略
        self.warmup_epochs = 5  # 学习率预热轮数
        self.warmup_lr = 0.0001  # 预热初始学习率
    
    def train_epoch(self, train_loader, epoch):
        """训练一个epoch，包含学习率预热"""
        # 学习率预热
        if epoch < self.warmup_epochs:
            # 线性预热
            warmup_factor = (epoch + 1) / self.warmup_epochs
            current_lr = self.warmup_lr + (self.optimizer.param_groups[0]['lr'] - self.warmup_lr) * warmup_factor
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = current_lr
            print(f"预热阶段，当前学习率: {current_lr:.6f}")
        
        return super().train_epoch(train_loader, epoch)


def create_resnet_config():
    """创建ResNet50的默认配置"""
    config = YAMLConfig()
    
    # 更新ResNet50专用配置
    resnet_config = {
        # 数据配置
        'data_dir': '../dataset/Mini-ImageNet-Dataset',
        'label_file': '../dataset/label.json',
        'batch_size': 256,  # ResNet可以使用较大的batch size
        'num_workers': 0,
        'pin_memory': True,
        'img_size': 224,
        'num_classes': 100,
        
        # 数据增强配置
        'normalize_mean': [0.485, 0.456, 0.406],
        'normalize_std': [0.229, 0.224, 0.225],
        
        # 训练配置
        'num_epochs': 120,  # ResNet通常需要更多epoch
        'lr': 0.1,  # ResNet使用较大的初始学习率
        'momentum': 0.9,
        'weight_decay': 1e-4,
        'print_freq': 50,
        
        # 学习率调度器
        'scheduler_type': 'StepLR',
        'step_size': 30,  # 每30个epoch学习率衰减
        'gamma': 0.1,
        
        # 路径配置
        'log_dir': './logs/resnet50',
        'model_save_path': './model/resnet50_model.pth',
        'save_best': True,
        
        # 模型配置
        'model_name': 'ResNet50',
        'model_type': 'resnet50',  # 可选: resnet18, resnet34, resnet50, resnet101, resnet152
        'dropout': 0.0,  # ResNet通常不使用dropout
        'use_warmup': True,  # 是否使用学习率预热
    }
    
    config.update(resnet_config)
    return config


def main():
    print("=== ResNet50 训练 ===")
    
    # 1. 创建配置
    config = create_resnet_config()
    
    # 检查是否存在YAML配置文件，如果存在则加载
    config_file = os.path.join(os.path.dirname(__file__), 'config.yaml')
    if os.path.exists(config_file):
        print(f"发现配置文件 {config_file}，加载中...")
        config = YAMLConfig(config_file)
    else:
        print("使用默认配置")
        # 保存默认配置到文件
        config.save_config(config_file)
        print(f"默认配置已保存到 {config_file}")
    
    # 打印配置信息
    print("\n=== 当前配置 ===")
    print(f"模型: {config.get('model_name')} ({config.get('model_type')})")
    print(f"批次大小: {config.get('batch_size')}")
    print(f"学习率: {config.get('lr')}")
    print(f"训练轮数: {config.get('num_epochs')}")
    print(f"图像尺寸: {config.get('img_size')}")
    print(f"Dropout: {config.get('dropout')}")
    print(f"学习率预热: {config.get('use_warmup')}")
    print("=" * 20)
    
    # 2. 设置设备
    device = setup_device()
    
    # 3. 创建数据加载器
    print("\n加载数据集...")
    train_loader, val_loader, test_loader, train_dataset = create_data_loaders(
        data_dir=config.get('data_dir'),
        batch_size=config.get('batch_size'),
        num_workers=config.get('num_workers'),
        pin_memory=config.get('pin_memory'),
        img_size=config.get('img_size'),
        mean=config.get('normalize_mean'),
        std=config.get('normalize_std')
    )
    
    print(f"训练集大小: {len(train_dataset)}")
    print(f"训练批次数: {len(train_loader)}")
    print(f"验证批次数: {len(val_loader)}")
    print(f"测试批次数: {len(test_loader)}")
    
    # 4. 加载标签映射
    label_mapping = load_label_mapping(config.get('label_file'))
    
    # 5. 创建模型
    print(f"\n创建{config.get('model_type', 'ResNet50')}模型...")
    
    # 根据配置选择模型类型
    model_type = config.get('model_type', 'resnet50').lower()
    if model_type == 'resnet18':
        model = ResNet18(num_classes=config.get('num_classes'))
    elif model_type == 'resnet34':
        model = ResNet34(num_classes=config.get('num_classes'))
    elif model_type == 'resnet50':
        model = ResNet50(num_classes=config.get('num_classes'), dropout=config.get('dropout'))
    elif model_type == 'resnet101':
        model = ResNet101(num_classes=config.get('num_classes'))
    elif model_type == 'resnet152':
        model = ResNet152(num_classes=config.get('num_classes'))
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")
    
    # 6. 创建损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(
        model.parameters(), 
        lr=config.get('lr'), 
        momentum=config.get('momentum'), 
        weight_decay=config.get('weight_decay')
    )
    
    # 7. 创建学习率调度器
    scheduler_type = config.get('scheduler_type', 'StepLR')
    if scheduler_type == 'StepLR':
        scheduler = optim.lr_scheduler.StepLR(
            optimizer, 
            step_size=config.get('step_size'), 
            gamma=config.get('gamma')
        )
    elif scheduler_type == 'CosineAnnealingLR':
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=config.get('T_max', 200)
        )
    else:
        scheduler = None
    
    # 8. 创建TensorBoard记录器
    tensorboard_logger = TensorBoardLogger(
        log_dir=config.get('log_dir'),
        model_name=config.get('model_name', 'resnet50'),
        comment='residual_network'
    )
    
    # 9. 创建训练器
    if config.get('use_warmup', True):
        trainer = ResNetTrainer(
            model=model,
            device=device,
            criterion=criterion,
            optimizer=optimizer,
            scheduler=scheduler,
            tensorboard_logger=tensorboard_logger,
            model_save_path=config.get('model_save_path'),
            print_freq=config.get('print_freq')
        )
    else:
        trainer = Trainer(
            model=model,
            device=device,
            criterion=criterion,
            optimizer=optimizer,
            scheduler=scheduler,
            tensorboard_logger=tensorboard_logger,
            model_save_path=config.get('model_save_path'),
            print_freq=config.get('print_freq')
        )
    
    # 10. 开始训练
    training_history = trainer.train(
        train_loader=train_loader,
        val_loader=val_loader,
        num_epochs=config.get('num_epochs'),
        save_best=config.get('save_best')
    )
    
    # 11. 测试模型
    print('\n加载最佳模型进行测试...')
    best_model_path = config.get('model_save_path').replace('.pth', '_best.pth')
    model = load_model(model, best_model_path, device)
    
    # 重新创建测试用的训练器
    test_trainer = Trainer(model=model, device=device)
    test_accuracy = test_trainer.test(
        test_loader=test_loader,
        train_dataset=train_dataset,
        label_mapping=label_mapping,
        num_classes=config.get('num_classes')
    )
    
    # 12. 记录最终结果
    tensorboard_logger.log_test_results(test_accuracy)
    tensorboard_logger.log_final_summary(
        training_time=training_history['training_time'],
        best_val_acc=training_history['best_val_acc'],
        test_acc=test_accuracy,
        model_path=config.get('model_save_path')
    )
    
    # 13. 保存最终配置
    final_config_path = os.path.join(config.get('log_dir'), 'final_config.yaml')
    config.save_config(final_config_path)
    print(f"\n最终配置已保存至: {final_config_path}")
    
    # 14. 打印最终结果
    print(f'\n=== ResNet50 训练总结 ===')
    print(f'训练时间: {training_history["training_time"]:.2f} 秒')
    print(f'最佳验证准确率: {training_history["best_val_acc"]:.2f}%')
    print(f'最终测试准确率: {test_accuracy:.2f}%')
    print(f'模型已保存至: {config.get("model_save_path")}')
    print('=' * 30)
    
    # 15. 关闭TensorBoard并打印命令
    tensorboard_logger.print_tensorboard_command()
    tensorboard_logger.close()


if __name__ == '__main__':
    main()
